Convert BST to greater tree

Time: O(N); Space: O(H); easy

Given a Binary Search Tree (BST), convert it to a Greater Tree such that:

every key of the original BST is changed to the original key plus sum of all keys greater than the original key in BST.

Example 1:

  5
 / \
2   13

Input: root = {TreeNode} [5,2,13]

  18
 /  \
20  13

Output: {TreeNode} [18,20,13]

Example 2:

  5
 / \
3  15

Input: root = {TreeNode} [5,3,15]

  20
 /  \
23  15

Output: {TreeNode} [20,23,15]

Notes:

[5]:

class TreeNode(object):
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

Auxiliary Tools

[6]:

from graphviz import Graph

class TreeTasks(object):
    def visualize_tree(self, tree):
        def add_nodes_edges(tree, dot=None):
            # Create Graph (not Digraph) object
            if dot is None:
                dot = Graph()
                dot.node(name=str(tree), label=str(tree.val))
            # Add nodes
            if tree.left:
                dot.node(name=str(tree.left), label="."+str(tree.left.val))
                dot.edge(str(tree), str(tree.left))
                dot = add_nodes_edges(tree.left, dot=dot)
            if tree.right:
                dot.node(name=str(tree.right), label=str(tree.right.val)+".")
                dot.edge(str(tree), str(tree.right))
                dot = add_nodes_edges(tree.right, dot=dot)
            return dot
        # Add nodes recursively and create a list of edges
        dot = add_nodes_edges(tree)
        # Visualize the graph
        display(dot)
        return dot

1. Recursion

[7]:

class Solution1(object):
    """
    Time: O(N)
    Space: O(H)
    """
    def convertBST(self, root):
        """
        :type root: TreeNode
        :rtype: TreeNode
        """
        def convertBSTHelper(root, cur_sum):
            if not root:
                return cur_sum

            if root.right:
                cur_sum = convertBSTHelper(root.right, cur_sum)
            cur_sum += root.val
            root.val = cur_sum

            if root.left:
                cur_sum = convertBSTHelper(root.left, cur_sum)
            return cur_sum

        convertBSTHelper(root, 0)

        return root

[8]:

s = Solution1()

root = TreeNode(5)
root.left, root.right = TreeNode(2), TreeNode(13)
tree = s.convertBST(root)
t = TreeTasks()
dot = t.visualize_tree(tree)
../../_images/topics_tree_0538_convert_bst_to_greater_tree_[O(N),O(H),easy]_6_0.svg
[9]:

root = TreeNode(5)
root.left, root.right = TreeNode(3), TreeNode(15)
tree = s.convertBST(root)
t = TreeTasks()
dot = t.visualize_tree(tree)
../../_images/topics_tree_0538_convert_bst_to_greater_tree_[O(N),O(H),easy]_7_0.svg